
##! ~/anaconda3/bin/python
# -*- coding: utf-8 -*-
# @Time   : 2019-11-15
# @Author : chujie
# @Email  : zjchenchujie@126.com
# @Brief  : find the annotations that unmatch image.

"""
This script is to check wether the images matches
annotations in PASCAL_VOC dataset.

"""

import os
import shutil
import argparse


def usage():
    """
    print usage information.
    :return:
    """
    print("[-] Usage: python voc_image_annotation_check.py --data_dir [VOC root dir]")


def check_dir(data_dir):
    """
    check images vs annotations,
    move unmatched files to data_dir/unmatched directory.
    :param data_dir: dataset directory
    :return: None
    """

    unmatched_xml_count = 0
    unmatched_img_count = 0
    display_flag = 0

    image_dir = os.path.join(data_dir, "JPEGImages/")
    xml_dir = os.path.join(data_dir, "Annotations/")
    unmatch_image_dir = os.path.join(data_dir, "unmatched/JPEGImages/")
    unmatch_xml_dir = os.path.join(data_dir, "unmatched/Annotations/")

    if not os.path.exists(image_dir):
        print("\"{}\" not exists. ".format(image_dir))
        exit("Please input a valid path.")
    if not os.path.exists(xml_dir):
        print("\"{}\" not exists. ".format(xml_dir))
        exit("Please input a valid path.")

    if not os.path.exists(unmatch_image_dir):
        os.makedirs(unmatch_image_dir)
    if not os.path.exists(unmatch_xml_dir):
        os.makedirs(unmatch_xml_dir)

    image_files = os.listdir(image_dir)
    img_file_num = len(image_files)

    print("==========Start, based on images.=========")
    for img_idx in range(img_file_num):
        display_flag += 1
        img_file_name = image_files[img_idx].split('.')[0]
        img_file_full_path = os.path.join(image_dir, image_files[img_idx])
        xml_file_full_path = os.path.join(xml_dir, "{}.xml".format(img_file_name))
        if not os.path.exists(xml_file_full_path):
            unmatched_xml_count += 1
            print("[-] Found \"{}\" unmatch. ".format(xml_file_full_path))
            shutil.move(img_file_full_path, os.path.join(unmatch_image_dir, image_files[img_idx]))
        if display_flag % 100 == 0:
            print("Processed {}, total {}".format(display_flag, img_file_num))
    print()
    display_flag = 0

    xml_files = os.listdir(xml_dir)
    xml_file_num = len(xml_files)
    print("==========Start, based on xmls.=========")
    for xml_idx in range(xml_file_num):
        display_flag += 1
        xml_file_name = xml_files[xml_idx].split('.')[0]
        xml_file_full_path = os.path.join(xml_dir, xml_files[xml_idx])
        img_file_full_path = os.path.join(image_dir, "{}.jpg".format(xml_file_name))
        if not os.path.exists(img_file_full_path):
            unmatched_img_count += 1
            print("[-] Found \"{}\" unmatch. ".format(img_file_full_path))
            shutil.move(xml_file_full_path, os.path.join(unmatch_xml_dir, xml_files[xml_idx]))
        if display_flag % 100 == 0:
            print("Processed {}, total {}".format(display_flag, img_file_num))
    print("Done!")
    print("Found {} unmatched image files. Found {} umatched xml files".
          format(unmatched_img_count, unmatched_xml_count))


if __name__ == '__main__':
    PARSER = argparse.ArgumentParser(description="argument parser")
    PARSER.add_argument("--data_dir", type=str,
                        help="enter the dataset directory in PASCAL VOC format.")
    ARGS = PARSER.parse_args()
    ROOT_DIR = ARGS.data_dir
    check_dir(ROOT_DIR)
